/*
 * GeoSystem.cpp
 *
 *  Created on: 22 Aug 2011
 *      Author: Allan
 */

#include "GeoSystem.h"

// C++ includes
#include <algorithm>
#include <cmath>
#include <iostream>
#include <iomanip>

// Boost includes
#include <boost/foreach.hpp>

// GeoReact includes
#include <Geochemistry/ElectrolyteSolution.h>
#include <Geochemistry/GeoDatabase.h>
#include <Utils/Algorithms.h>

GeoSystem::GeoSystem()
{
	SetState(25.0, 1.0, 1.0E-7);
}

GeoSystem::GeoSystem(
	const Multiphase& multiphase, 
	const ReactionSystem& equilibrium_reactions)
: multiphase(multiphase), equilibrium_reactions(equilibrium_reactions)
{
	InitializeStoichiometricMatrixUki();
	InitializeMolecularWeights();
	SetState(25.0, 1.0, 1.0E-7);
}

GeoSystem::GeoSystem(
	const Multiphase& multiphase,
	const ReactionSystem& equilibrium_reactions,
	const ReactionSystem& kinetic_reactions,
	const vector<string>& primary_species, 
	const vector<string>& equilibrium_species,
	const vector<string>& kinetic_species)
: multiphase(multiphase), equilibrium_reactions(equilibrium_reactions),
  kinetic_reactions(kinetic_reactions), primary_species(primary_species),
  equilibrium_species(equilibrium_species), kinetic_species(kinetic_species)
{
	InitializeStoichiometricMatrixUki();
	InitializeMolecularWeights();
	SetState(25.0, 1.0, 1.0E-7);
}

void
GeoSystem::SetTemperature(double T)
{
	SetState(T, this->P, this->n);
}

void
GeoSystem::SetPressure(double P)
{
	SetState(this->T, P, this->n);
}

void
GeoSystem::SetComposition(double n)
{
	SetState(this->T, this->P, n);
}

void
GeoSystem::SetComposition(const VectorXd& n)
{
	SetState(this->T, this->P, n);
}

void
GeoSystem::SetState(double T, double P)
{
	SetState(T, P, this->n);
}

void
GeoSystem::SetState(double T, double P, double n)
{
	SetState(T, P, VectorXd::Constant(multiphase.GetNumSpecies(), n));
}

void
GeoSystem::SetState(double T, double P, const VectorXd& n)
{
	this->T = T;
	this->P = P;
	this->n = n;
	
	a = multiphase.Activities(T, P, n);
	g = multiphase.ActivityCoefficients(T, P, n);
	
	if(kinetic_reactions.GetNumReactions()) 
		r = uki.transpose() * kinetic_reactions.KineticRates(T, P, n, a);
	else
		r = VectorXd::Zero(multiphase.GetNumSpecies());
}

const double
GeoSystem::GetTemperature() const
{
	return T;
}

const double
GeoSystem::GetPressure() const
{
	return P;
}

const VectorXd&
GeoSystem::GetComposition() const
{
	return n;
}

const VectorXd&
GeoSystem::GetActivities() const
{
	return a;
}

const VectorXd&
GeoSystem::GetActivityCoefficients() const
{
	return g;
}

const VectorXd&
GeoSystem::GetKineticRates() const
{
	return r;
}

const Multiphase&
GeoSystem::GetMultiphase() const
{
	return multiphase;
}

const ReactionSystem&
GeoSystem::GetEquilibriumReactions() const
{
	return equilibrium_reactions;
}

const ReactionSystem&
GeoSystem::GetKineticReactions() const
{
	return kinetic_reactions;
}

const vector<string>&
GeoSystem::GetPrimarySpecies() const
{
	return primary_species;
}

const vector<string>&
GeoSystem::GetEquilibriumSpecies() const
{
	return equilibrium_species;
}

const vector<string>&
GeoSystem::GetKineticSpecies() const
{
	return kinetic_species;
}

const double
GeoSystem::Moles(const string& species) const
{
	return n[multiphase[species]];
}

const double
GeoSystem::Mass(const string& species) const
{
	return Moles(species) * molecular_weights[multiphase[species]];
}

const double
GeoSystem::Activity(const string& species) const
{
	return a[multiphase[species]];
}

const double
GeoSystem::ActivityCoefficient(const string& species) const
{
	return g[multiphase[species]];
}

const double
GeoSystem::KineticRate(const string& species) const
{
	return r[multiphase[species]];
}

const double
GeoSystem::Acidity() const
{
	return -log10(a[multiphase["H[+]"]]);
}

const double
GeoSystem::IonicStrengthEffective() const
{
	const Index index_aqueous_phase = multiphase.FindPhase("Aqueous");
	
	const Phase& aqueous_phase = multiphase.GetPhases()[index_aqueous_phase];
	
	const vector<string>& aqueous_species = aqueous_phase.GetSpecies();
	
	ElectrolyteSolution electrolyte_solution(aqueous_species);

	const VectorXd n_aqueous_phase = multiphase.PhaseMoles(index_aqueous_phase, n);
	
	return electrolyte_solution.EffectiveIonicStrength(n_aqueous_phase);
}

const double
GeoSystem::IonicStrengthStoichiometric() const
{
	const Index index_aqueous_phase = multiphase.FindPhase("Aqueous");
	
	const Phase& aqueous_phase = multiphase.GetPhases()[index_aqueous_phase];
	
	const vector<string>& aqueous_species = aqueous_phase.GetSpecies();
	
	ElectrolyteSolution electrolyte_solution(aqueous_species);

	const VectorXd n_aqueous_phase = multiphase.PhaseMoles(index_aqueous_phase, n);
	
	return electrolyte_solution.StoichiometricIonicStrength(n_aqueous_phase);
}

const double
GeoSystem::PartialPressureGas(const string& gas) const
{
	const Index index_gas = multiphase[gas];

	const Index index_gaseous_phase = multiphase.FindPhase("Gaseous");
	
	const VectorXd n_gaseous_phase = multiphase.PhaseMoles(index_gaseous_phase, n);
	
	return n[index_gas] / n_gaseous_phase.sum() * P;
}

void
GeoSystem::InitializeStoichiometricMatrixUki()
{
	if(kinetic_reactions.GetNumReactions())
		uki = kinetic_reactions.AssembleStoichiometricMatrix(multiphase.GetSpecies());
}

void
GeoSystem::InitializeMolecularWeights()
{
	molecular_weights.resize(multiphase.GetNumSpecies());
	
	BOOST_FOREACH(const string& species, multiphase.GetSpecies())
		molecular_weights[multiphase[species]] = GeoDatabase::MolecularWeight(species);
}

void OutputMultiphaseState(const GeoSystem& geosystem, std::ostream& out)
{
	out << ">>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<" << endl;
	out << ">>>            Multiphase            <<<" << endl;
	out << ">>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<" << endl;
	
	vector<string> species = geosystem.GetMultiphase().GetSpecies();
	
	auto compare = [&](const string& s1, const string& s2)
		{ return geosystem.Moles(s1) > geosystem.Moles(s2); };
	
	sort(species.begin(), species.end(), compare);
	
	out << setprecision(6);
	
	out << setw(25) << left << "Species";
	out << setw(25) << left << "Moles";
	out << setw(25) << left << "Mass";
	out << setw(25) << left << "Activity";
	out << setw(25) << left << "ActivityCoefficient" << endl;
	
	BOOST_FOREACH(const string& s, species)
	{
		out << setw(25) << left << s;
		out << setw(25) << left << geosystem.Moles(s);
		out << setw(25) << left << geosystem.Mass(s);
		out << setw(25) << left << geosystem.Activity(s);
		out << setw(25) << left << geosystem.ActivityCoefficient(s) << endl;
	}
}

void OutputReactionsState(const GeoSystem& geosystem, std::ostream& out)
{
	const ReactionSystem& equilibrium_reactions = geosystem.GetEquilibriumReactions();
	const ReactionSystem& kinetic_reactions = geosystem.GetKineticReactions();
	
	double T = geosystem.GetTemperature();
	double P = geosystem.GetPressure();
	
	const VectorXd& n = geosystem.GetComposition();
	const VectorXd& a = geosystem.GetActivities();
	
	auto output_reaction = [&](const Reaction& reaction)
		{
			const double K = reaction.EquilibriumConst(T, P);
			const double Q = reaction.ReactionQuotient(a);
			const double kinetic_rate = reaction.KineticRate(T, P, n, a);
			
			out << "Reaction: " << reaction << endl;
			out << "      log(K) = " << log(K) << endl;
			out << "      log(Q) = " << log(Q) << endl;
			out << "    log(Q/K) = " << log(Q/K) << endl;
			out << "        Rate = " << kinetic_rate << " mol/s" << endl << endl;
		};
	
	if(equilibrium_reactions.GetNumReactions() > 0)
	{
		out << ">>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<" << endl;
		out << ">>> Equilibrium-Controlled Reactions <<<" << endl;
		out << ">>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<" << endl;
		
		BOOST_FOREACH(auto reaction, equilibrium_reactions.GetReactions())
			output_reaction(reaction);
	}
	
	if(kinetic_reactions.GetNumReactions() > 0)
	{
		out << ">>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<" << endl;
		out << ">>> Kinetically-Controlled Reactions <<<" << endl;
		out << ">>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<" << endl;
		
		BOOST_FOREACH(auto reaction, kinetic_reactions.GetReactions())
			output_reaction(reaction);
	}
}

std::ostream& operator<<(std::ostream& out, const GeoSystem& geosystem)
{
	OutputMultiphaseState(geosystem, out);
	
	out << endl;
	
	OutputReactionsState(geosystem, out);
	
	return out;
}
